{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "174ef1d9",
   "metadata": {},
   "source": [
    "# Parent Document Retriever\n",
    "\n",
    "**문서 검색과 문서 분할의 균형 잡기**\n",
    "\n",
    "문서 검색 과정에서 문서를 적절한 크기의 조각(청크)으로 나누는 것은 다음의 **상충되는 두 가지 중요한 요소를 고려**해야 합니다.\n",
    "\n",
    "1. 작은 문서를 원하는 경우: 이렇게 하면 문서의 임베딩이 그 의미를 가장 정확하게 반영할 수 있습니다. 문서가 너무 길면 임베딩이 의미를 잃어버릴 수 있습니다.\n",
    "2. 각 청크의 맥락이 유지되도록 충분히 긴 문서를 원하는 경우입니다.\n",
    "\n",
    "**`ParentDocumentRetriever`의 역할**\n",
    "\n",
    "이 두 요구 사항 사이의 균형을 맞추기 위해 `ParentDocumentRetriever`라는 도구가 사용됩니다. 이 도구는 문서를 작은 조각으로 나누고, 이 조각들을 관리합니다. 검색을 진행할 때는, 먼저 이 작은 조각들을 찾아낸 다음, 이 조각들이 속한 원본 문서(또는 더 큰 조각)의 식별자(ID)를 통해 전체적인 맥락을 파악할 수 있습니다.\n",
    "\n",
    "여기서 '부모 문서'란, 작은 조각이 나누어진 원본 문서를 말합니다. 이는 전체 문서일 수도 있고, 비교적 큰 다른 조각일 수도 있습니다. 이 방식을 통해 문서의 의미를 정확하게 파악하면서도, 전체적인 맥락을 유지할 수 있게 됩니다.\n",
    "\n",
    "**정리**\n",
    "\n",
    "- **문서 간의 계층 구조 활용**: `ParentDocumentRetriever`는 문서 검색의 효율성을 높이기 위해 문서 간의 계층 구조를 활용합니다.\n",
    "- **검색 성능 향상**: 관련성 높은 문서를 빠르게 찾아내며, 주어진 질문에 대한 가장 적합한 답변을 제공하는 문서를 효과적으로 찾아낼 수 있습니다.\n",
    "  문서를 검색할 때 자주 발생하는 두 가지 상충되는 요구 사항이 있습니다:\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7ba411b",
   "metadata": {},
   "source": [
    "여러 개의 텍스트 파일을 로드하기 위해 `TextLoader` 객체를 생성하고 데이터를 로드합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d14d13cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "514e4b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH10-Retriever\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c12ba1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.storage import InMemoryStore\n",
    "from langchain_community.document_loaders import TextLoader\n",
    "from langchain_chroma import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain.retrievers import ParentDocumentRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68f6a054",
   "metadata": {},
   "outputs": [],
   "source": [
    "loaders = [\n",
    "    # 파일을 로드합니다.\n",
    "    TextLoader(\"./data/appendix-keywords.txt\"),\n",
    "]\n",
    "\n",
    "docs = []\n",
    "for loader in loaders:\n",
    "    # 로더를 사용하여 문서를 로드하고 docs 리스트에 추가합니다.\n",
    "    docs.extend(loader.load())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5314377a",
   "metadata": {},
   "source": [
    "## 전체 문서 검색\n",
    "\n",
    "이 모드에서는 전체 문서를 검색하고자 합니다. 따라서 `child_splitter` 만 지정하도록 하겠습니다.\n",
    "\n",
    "- 나중에는 `parent_splitter` 도 지정하여 결과를 비교해 보겠습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "812dfdfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 자식 분할기를 생성합니다.\n",
    "child_splitter = RecursiveCharacterTextSplitter(chunk_size=200)\n",
    "\n",
    "# DB를 생성합니다.\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"full_documents\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "\n",
    "store = InMemoryStore()\n",
    "\n",
    "# Retriever 를 생성합니다.\n",
    "retriever = ParentDocumentRetriever(\n",
    "    vectorstore=vectorstore,\n",
    "    docstore=store,\n",
    "    child_splitter=child_splitter,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02d1967f",
   "metadata": {},
   "source": [
    "`retriever.add_documents(docs, ids=None)` 함수로 문서목록을 추가합니다.\n",
    "\n",
    "- `ids` 가 `None` 이면 자동으로 생성됩니다.\n",
    "- `add_to_docstore=False` 로 설정시 document 를 중복으로 추가하지 않습니다. 단, 중복을 체크하기 위한 `ids` 값이 필수 값으로 요구됩니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c737c91d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서를 검색기에 추가합니다. docs는 문서 목록이고, ids는 문서의 고유 식별자 목록입니다.\n",
    "retriever.add_documents(docs, ids=None, add_to_docstore=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68ef1793",
   "metadata": {},
   "source": [
    "이 코드는 두 개의 키를 반환해야 합니다. 그 이유는 우리가 두 개의 문서를 추가했기 때문입니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72bc0f5b",
   "metadata": {},
   "source": [
    "- `store` 객체의 `yield_keys()` 메서드를 호출하여 반환된 키(key) 값들을 리스트로 변환합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6875b63a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 저장소의 모든 키를 리스트로 반환합니다.\n",
    "list(store.yield_keys())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9fb5f346",
   "metadata": {},
   "source": [
    "이제 벡터 스토어 검색 기능을 호출해 보겠습니다.\n",
    "\n",
    "우리가 작은 청크(chunk)들을 저장하고 있기 때문에, 검색 결과로 작은 청크들이 반환되는 것을 확인할 수 있을 것입니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c15a995",
   "metadata": {},
   "source": [
    "`vectorstore` 객체의 `similarity_search` 메서드를 사용하여 유사도 검색을 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6ed9a29",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 유사도 검색을 수행합니다.\n",
    "sub_docs = vectorstore.similarity_search(\"Word2Vec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f377d0d9",
   "metadata": {},
   "source": [
    "`sub_docs[0].page_content`를 출력합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4db52fa3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# sub_docs 리스트의 첫 번째 요소의 page_content 속성을 출력합니다.\n",
    "print(sub_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "399e28a3",
   "metadata": {},
   "source": [
    "이제 전체 retriever에서 검색해 보겠습니다. 이 과정에서는 작은 청크(chunk)들이 위치한 **문서를 반환** 하기 때문에 상대적으로 큰 문서들이 반환될 것입니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f21c8dce",
   "metadata": {},
   "source": [
    "`retriever` 객체의 `invoke()` 메서드를 사용하여 쿼리와 관련된 문서를 검색합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd845482",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서를 검색하여 가져옵니다.\n",
    "retrieved_docs = retriever.invoke(\"Word2Vec\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a3cb5b8",
   "metadata": {},
   "source": [
    "검색된 문서(`retrieved_docs[0]`)의 일부 내용을 출력합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "531585d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 검색된 문서의 문서의 페이지 내용의 길이를 출력합니다.\n",
    "print(\n",
    "    f\"문서의 길이: {len(retrieved_docs[0].page_content)}\",\n",
    "    end=\"\\n\\n=====================\\n\\n\",\n",
    ")\n",
    "\n",
    "# 문서의 일부를 출력합니다.\n",
    "print(retrieved_docs[0].page_content[2000:2500])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "829387f8",
   "metadata": {},
   "source": [
    "## 더 큰 Chunk 의 크기를 조절\n",
    "\n",
    "이전의 결과처럼 **전체 문서가 너무 커서 있는 그대로 검색하기에는 부적합** 할 수 있습니다.\n",
    "\n",
    "이런 경우, 실제로 우리가 하고 싶은 것은 먼저 원시 문서를 더 큰 청크로 분할한 다음, 더 작은 청크로 분할하는 것입니다.\n",
    "\n",
    "그런 다음 작은 청크들을 인덱싱하지만, 검색 시에는 더 큰 청크를 검색합니다 (그러나 여전히 전체 문서는 아닙니다).\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "51dcc001",
   "metadata": {},
   "source": [
    "- `RecursiveCharacterTextSplitter`를 사용하여 부모 문서와 자식 문서를 생성합니다.\n",
    "  - 부모 문서는 `chunk_size`가 1000으로 설정되어 있습니다.\n",
    "  - 자식 문서는 `chunk_size`가 200으로 설정되어 있으며, 부모 문서보다 작은 크기로 생성됩니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29662c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 부모 문서를 생성하는 데 사용되는 텍스트 분할기입니다.\n",
    "parent_splitter = RecursiveCharacterTextSplitter(chunk_size=1000)\n",
    "# 자식 문서를 생성하는 데 사용되는 텍스트 분할기입니다.\n",
    "# 부모보다 작은 문서를 생성해야 합니다.\n",
    "child_splitter = RecursiveCharacterTextSplitter(chunk_size=200)\n",
    "# 자식 청크를 인덱싱하는 데 사용할 벡터 저장소입니다.\n",
    "vectorstore = Chroma(\n",
    "    collection_name=\"split_parents\", embedding_function=OpenAIEmbeddings()\n",
    ")\n",
    "# 부모 문서의 저장 계층입니다.\n",
    "store = InMemoryStore()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db4a3592",
   "metadata": {},
   "source": [
    "`ParentDocumentRetriever`를 초기화하는 코드입니다.\n",
    "\n",
    "- `vectorstore` 매개변수는 문서 벡터를 저장하는 벡터 저장소를 지정합니다.\n",
    "- `docstore` 매개변수는 문서 데이터를 저장하는 문서 저장소를 지정합니다.\n",
    "- `child_splitter` 매개변수는 하위 문서를 분할하는 데 사용되는 문서 분할기를 지정합니다.\n",
    "- `parent_splitter` 매개변수는 상위 문서를 분할하는 데 사용되는 문서 분할기를 지정합니다.\n",
    "\n",
    "`ParentDocumentRetriever`는 계층적 문서 구조를 처리하며, 상위 문서와 하위 문서를 별도로 분할하고 저장합니다. 이를 통해 검색 시 상위 문서와 하위 문서를 효과적으로 활용할 수 있습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3cacb06",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever = ParentDocumentRetriever(\n",
    "    # 벡터 저장소를 지정합니다.\n",
    "    vectorstore=vectorstore,\n",
    "    # 문서 저장소를 지정합니다.\n",
    "    docstore=store,\n",
    "    # 하위 문서 분할기를 지정합니다.\n",
    "    child_splitter=child_splitter,\n",
    "    # 상위 문서 분할기를 지정합니다.\n",
    "    parent_splitter=parent_splitter,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc96735d",
   "metadata": {},
   "source": [
    "`retriever` 객체에 `docs`를 추가합니다. `retriever`가 검색할 수 있는 문서 집합에 새로운 문서들을 추가하는 역할을 합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f95883b",
   "metadata": {},
   "outputs": [],
   "source": [
    "retriever.add_documents(docs)  # 문서를 retriever에 추가합니다."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0a560f6",
   "metadata": {},
   "source": [
    "이제 문서의 수가 훨씬 더 많아진 것을 볼 수 있습니다. 이는 더 큰 청크(chunk)들입니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d0989cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 저장소에서 키를 생성하고 리스트로 변환한 후 길이를 반환합니다.\n",
    "len(list(store.yield_keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8004388",
   "metadata": {},
   "source": [
    "기본 벡터 저장소가 여전히 작은 청크를 검색하는지 확인해 보겠습니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6ea50f0",
   "metadata": {},
   "source": [
    "`vectorstore` 객체의 `similarity_search` 메서드를 사용하여 유사도 검색을 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec52ff30",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 유사도 검색을 수행합니다.\n",
    "sub_docs = vectorstore.similarity_search(\"Word2Vec\")\n",
    "# sub_docs 리스트의 첫 번째 요소의 page_content 속성을 출력합니다.\n",
    "print(sub_docs[0].page_content)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "937b2c64",
   "metadata": {},
   "source": [
    "이번에는 `retriever` 객체의 `invoke()` 메서드를 사용하여 문서를 검색합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9944a690",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서를 검색하여 가져옵니다.\n",
    "retrieved_docs = retriever.invoke(\"Word2Vec\")\n",
    "\n",
    "# 검색된 문서의 첫 번째 문서의 페이지 내용의 길이를 반환합니다.\n",
    "print(retrieved_docs[0].page_content)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
